#pragma once

#include "common_unit.h"

class DebugUnit{
public:
    static const char* device_type_str[3];
public:
    static void mem_init(std::shared_ptr<MemoryNode> mem_node, void * arg_addr = nullptr) {
        void * addr = mem_node->storage_ptr;
        if (arg_addr) addr = arg_addr;
        size_t size = mem_node->total_size;
        // size = (size + 63) & (~63);
        int num_ele = (int)size / (int)sizeof(uint64_t);
        for (int i = 0; i < num_ele; ++i) {
            if (mem_node->next_opt_nodes[0]->op_type == "double") {
                ( (double *) addr)[i] = (double)i;
            }
            else {
                ( (int64_t * ) addr)[i] = i - 3;
            }
        }
    }

    static void mem_debug(std::shared_ptr<MemoryNode> mem_node, int limit=-1, void * arg_addr = nullptr) {
        printf ("mem feature = : %s\n", mem_node->name.c_str() );
        void * addr = mem_node->storage_ptr;
        if (arg_addr) addr = arg_addr;
        int feature_num = mem_node->feature_end_id - mem_node->feature_start_id;
        int batch_size = mem_node->batch_end_id - mem_node->batch_start_id;
        int num_ele_per_feature = mem_node->element_per_feature;

        for (int i = 0; i < feature_num; ++i) {
            printf ("(");
            int Max_printf = batch_size;
            if (limit != -1) Max_printf = std::min(Max_printf, limit);
            for (int j = 0; j < Max_printf; ++j) {
                printf ("[");
                for (int k = 0; k < num_ele_per_feature; ++k) {
                    if ( (int)mem_node->prev_opt_nodes.size() > 0 && mem_node->prev_opt_nodes[0]->op_type == "embedding_fetch") {
                        printf ("<");
                        for (int o = 0; o < 10; ++o) {
                            float * pos = (float *) ( (uint64_t) addr + 
                            (i * batch_size * num_ele_per_feature * EMBEDDING_DIM + 
                                j * num_ele_per_feature * EMBEDDING_DIM + k * EMBEDDING_DIM + o) * sizeof(float) );
                            printf ("%lf ", *((float *)pos) );
                        } printf (">");
                    }
                    else {
                        int64_t * pos = (int64_t *) ( (uint64_t) addr + 
                        (i * batch_size * num_ele_per_feature + j * num_ele_per_feature + k) * sizeof(uint64_t) );
                        
                        if (mem_node->dtype == "double") printf ("%lf ", *((double *)pos) );
                        else printf ("%ld ", * ( (volatile uint64_t*)pos) );
                    }
                }
                printf ("]");
            }
            printf (")\n");
        }
    }

    static void operator_debug(std::shared_ptr<OperatorNode> op_node) {
        printf ("Input mem data:\n");
        for (auto mem_node : op_node->prev_mem_nodes) {
            DebugUnit::mem_debug(mem_node);
        }
        printf ("execute operator %s on device %s\n", op_node->op_type.c_str(), DebugUnit::device_type_str[op_node->device_type] );
        printf ("Output mem data:\n");
        for (auto mem_node : op_node->next_mem_nodes) {
            DebugUnit::mem_debug(mem_node);
        }
    }
};

const char* DebugUnit::device_type_str[3] = {
    "DEVICE_CPU",
    "DEVICE_ARM",
    "DEVICE_DPA"
};
